from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from shutil import move
from time import time
from types import SimpleNamespace
from typing import TYPE_CHECKING, List, Tuple

from requests import exceptions
from requests import get

from src.custom.static import DESCRIPTION_LENGTH
from src.custom.static import MAX_WORKERS
from src.tools.retry import PrivateRetry

if TYPE_CHECKING:
    from src.config.parameter import Parameter


class Downloader:
    Phone_headers = {
        'User-Agent': 'com.ss.android.ugc.trill/494+Mozilla/5.0+(Linux;+Android+12;+2112123G+Build/SKQ1.211006.001;+wv)'
                      '+AppleWebKit/537.36+(KHTML,+like+Gecko)+Version/4.0+Chrome/107.0.5304.105+Mobile+Safari/537.36'}

    def __init__(self, params: "Parameter"):
        self.cleaner = params.cleaner
        self.cookie = params.cookie_object
        self.PC_headers, self.black_headers = self.init_headers(params.headers)
        self.PC_headers_tiktok = params.headers_tiktok
        self.xb = params.xb
        self.proxies = params.proxies
        self.download = params.download
        self.max_size = params.max_size
        self.chunk = params.chunk
        self.max_retry = params.max_retry
        self.blacklist = params.blacklist
        self.timeout = params.timeout
        self.ffmpeg = params.ffmpeg
        self.__thread = ThreadPoolExecutor
        self.__pool = None
        self.__temp = params.temp
        self.split = params.split
        self.original_cover = params.original_cover
        self.dynamic_cover = params.dynamic_cover
        self.music = params.music
        self.folder_name = params.folder_name
        self.root = params.root
        self.name_format = params.name_format
        self.folder_mode = params.folder_mode

    @staticmethod
    def init_headers(headers: dict) -> tuple:
        return headers, {"User-Agent": headers["User-Agent"]}

    def run(self,
            data: List[dict],
            type_: str,
            **kwargs, ) -> None:
        if not self.download:
            return
        if type_=="batch":
            print("开始下载作品文件")
            self.run_batch(data, **kwargs)
        elif type_=="works":
            self.run_general(data, **kwargs)
        else:
            raise ValueError

    def run_batch(
            self,
            data: List[dict],
            id_: str,
            name: str,
            mark="",
            addition="发布作品",
            mid: str = None,
            title: str = None,
    ):
        # assert addition in {"喜欢作品", "收藏作品", "发布作品", "合集作品"}, ValueError
        mix = addition=="合集作品"
        root = self.storage_folder(
            mid if mix else id_,
            title if mix else name,
            not mix,
            mark,
            addition,
            mix)
        self.batch_processing(data, root)

    def run_general(self, data: List[dict], tiktok: bool):
        root = self.storage_folder()
        self.batch_processing(data, root, False, tiktok=tiktok)

    def run_live(self, data: List[tuple]):
        if not data or not self.download:
            return []
        download_tasks = []
        download_command = []
        processIds = []
        self.generate_live_tasks(data, download_tasks, download_command)
        if self.ffmpeg.state:
            print("检测到 ffmpeg，程序将会调用 ffmpeg 下载直播，关闭 TikTokDownloader 不会中断下载！")
            processIds = self.__download_live(download_command)
        else:
            print("未检测到 ffmpeg，程序将会调用内置下载器下载直播，您需要保持 TikTokDownloader 运行直到直播结束！")
            self.downloader_chart(
                download_tasks,
                SimpleNamespace(),
                len(download_tasks),
                unknown_size=True,
                headers=self.black_headers)
        return processIds

    def generate_live_tasks(
            self, data: List[tuple], tasks: list, commands: list):
        for i, f, m in data:
            name = self.cleaner.filter_name(
                f'{i["title"]}{self.split}{i["nickname"]}{self.split}{datetime.now():%Y-%m-%d %H.%M.%S}.flv',
                inquire=False,
                default=str(
                    time())[
                        :10])
            temp_root, actual_root = self.deal_folder_path(
                self.storage_folder(folder_name="Live"), name, True)
            tasks.append((
                f,
                temp_root,
                actual_root,
                f'直播 {i["title"]}{self.split}{i["nickname"]}',
                "0" * 19,
            ))
            commands.append((
                m,
                str(actual_root.with_name(f"{actual_root.stem}.mp4").resolve()),
            ))

    def __download_live(self, commands: list):
        return self.ffmpeg.download(
            commands,
            self.proxies["https"],
            self.timeout,
            self.black_headers["User-Agent"],
        )

    def batch_processing(
            self,
            data: List[dict],
            root: Path,
            statistics=True,
            **kwargs):
        count = SimpleNamespace(
            downloaded_image=set(),
            skipped_image=set(),
            downloaded_video=set(),
            skipped_video=set()
        )
        tasks = []
        for item in data:
            item["desc"] = item["desc"][:DESCRIPTION_LENGTH]
            name = self.generate_works_name(item)
            temp_root, actual_root = self.deal_folder_path(root, name)
            params = {
                "tasks": tasks,
                "name": name,
                "id_": item["id"],
                "item": item,
                "count": count,
                "temp_root": temp_root,
                "actual_root": actual_root
            }
            if (t := item["type"])=="图集":
                self.download_image(**params)
            elif t=="视频":
                self.download_video(**params)
            self.download_music(**params)
            self.download_cover(**params)
        self.downloader_chart(
            tasks,
            count,
            **kwargs)
        if statistics:
            self.statistics_count(count)

    def downloader_chart(
            self,
            tasks: List[tuple],
            count: SimpleNamespace,
            max_workers=MAX_WORKERS,
            **kwargs):
        with self.__thread(max_workers=max_workers) as self.__pool:
            futures = []
            for task in tasks:
                future = self.__pool.submit(
                    self.request_file,
                    *task,
                    count=count,
                    **kwargs)
                futures.append(future)

            for future in futures:
                if not future.result():
                    # If any task fails, cancel all other tasks
                    for future in futures:
                        future.cancel()
                    break

    def deal_folder_path(self, root: Path, name: str,
                         pass_=False) -> Tuple[Path, Path]:
        root = self.create_works_folder(root, name, pass_)
        root.mkdir(exist_ok=True)
        temp = self.__temp.joinpath(name)
        actual = root.joinpath(name)
        return temp, actual

    def is_in_blacklist(self, id_: str) -> bool:
        return id_ in self.blacklist.record

    @staticmethod
    def is_exists(path: Path) -> bool:
        return path.exists()

    def is_skip(self, id_: str, path: Path) -> bool:
        return self.is_in_blacklist(id_) or self.is_exists(path)

    def download_image(
            self,
            tasks: list,
            name: str,
            id_: str,
            item: SimpleNamespace,
            count: SimpleNamespace,
            temp_root: Path,
            actual_root: Path) -> None:
        for index, img in enumerate(
                item["downloads"].split(" "), start=1):
            if self.is_in_blacklist(id_):
                count.skipped_image.add(id_)
                print(f"图集 {id_} 存在下载记录，跳过下载")
                count.skipped_image.add(id_)
                break
            elif self.is_exists(p := actual_root.with_name(f"{name}_{index}.jpeg")):
                print(f"图集 {id_}_{index} 文件已存在，跳过下载")
                print(f"文件路径: {p.resolve()}")
                count.skipped_image.add(id_)
                continue
            tasks.append((
                img,
                temp_root.with_name(
                    f"{name}_{index}.jpeg"),
                p,
                f"图集 {id_}_{index}",
                id_,
            ))

    def download_video(
            self,
            tasks: list,
            name: str,
            id_: str,
            item: SimpleNamespace,
            count: SimpleNamespace,
            temp_root: Path,
            actual_root: Path) -> None:
        if self.is_skip(
                id_, p := actual_root.with_name(
                    f"{name}.mp4")):
            print(f"视频 {id_} 存在下载记录或文件已存在，跳过下载")
            print(f"文件路径: {p.resolve()}")
            count.skipped_video.add(id_)
            return
        tasks.append((
            item["downloads"],
            temp_root.with_name(f"{name}.mp4"),
            p,
            f"视频 {id_}",
            id_,
        ))

    def download_music(
            self,
            tasks: list,
            name: str,
            id_: str,
            item: SimpleNamespace,
            temp_root: Path,
            actual_root: Path) -> None:
        if self.check_deal_music(
                url := item["music_url"],
                p := actual_root.with_name(f"{name}.mp3"), ):
            tasks.append((
                url,
                temp_root.with_name(f"{name}.mp3"),
                p,
                f"音乐 {id_}",
                id_,
            ))

    def download_cover(
            self,
            tasks: list,
            name: str,
            id_: str,
            item: SimpleNamespace,
            temp_root: Path,
            actual_root: Path) -> None:
        if all((self.original_cover,
                url := item["origin_cover"],
                not self.is_exists(p := actual_root.with_name(f"{name}.jpeg"))
                )):
            tasks.append((
                url,
                temp_root.with_name(f"{name}.jpeg"),
                p,
                f"封面 {id_}",
                id_,
            ))
        if all((self.dynamic_cover,
                url := item["dynamic_cover"],
                not self.is_exists(p := actual_root.with_name(f"{name}.webp"))
                )):
            tasks.append((
                url,
                temp_root.with_name(f"{name}.webp"),
                p,
                f"动图 {id_}",
                id_,
            ))

    def check_deal_music(self, url: str, path: Path) -> bool:
        return all((self.music, url, not self.is_exists(path)))

    @PrivateRetry.retry
    def request_file(
            self,
            url: str,
            temp: Path,
            actual: Path,
            show: str,
            id_: str,
            count: SimpleNamespace,
            headers: dict = None,
            tiktok=False,
            unknown_size=False) -> bool:
        try:
            with get(
                    url,
                    stream=True,
                    proxies=self.proxies,
                    headers=self.__adapter_headers(headers, tiktok),
                    timeout=self.timeout) as response:
                if not (
                        content := int(
                            response.headers.get(
                                'content-length',
                                0))) and not unknown_size:
                    print(f"{url} 响应内容为空")
                    return False
                if response.status_code!=200:
                    print(f"{response.url} 响应码异常: {response.status_code}")
                    return False
                elif all((self.max_size, content, content > self.max_size)):
                    print(f"{show} 文件大小超出限制，跳过下载")
                    return True
                with open(actual, 'wb') as file:
                    for chunk in response.iter_content(chunk_size=1024):
                        if chunk:
                            file.write(chunk)
                return True
        except (exceptions.ConnectionError,
                exceptions.ChunkedEncodingError,
                exceptions.ReadTimeout) as e:
            print(f"网络异常: {e}")
            return False

    def download_file(
            self,
            temp: Path,
            actual: Path,
            show: str,
            id_: str,
            response,
            count: SimpleNamespace) -> bool:
        try:
            with temp.open("wb") as f:
                for chunk in response.iter_content(chunk_size=self.chunk):
                    f.write(chunk)
        except exceptions.ChunkedEncodingError:
            print(f"{show} 由于网络异常下载中断")
            self.delete_file(temp)
            return False
        self.save_file(temp, actual)
        print(f"{show} 文件下载成功")
        print(f"文件路径 {actual.resolve()}")
        self.blacklist.update_id(id_)
        self.add_count(show, id_, count)
        return True

    def __adapter_headers(
            self,
            headers: dict,
            tiktok: bool) -> dict:
        return headers or (
            self.Phone_headers if tiktok else self.black_headers)

    @staticmethod
    def add_count(type_: str, id_: str, count: SimpleNamespace):
        if type_.startswith("图集"):
            count.downloaded_image.add(id_)
        elif type_.startswith("视频"):
            count.downloaded_video.add(id_)

    def storage_folder(
            self,
            id_: str = None,
            name: str = None,
            batch=False,
            mark: str = None,
            addition: str = None,
            mix=False,
            folder_name: str = None) -> Path:
        if batch and all((id_, name, addition)):
            folder_name = f"UID{id_}_{mark or name}_{addition}"
        elif mix and all((id_, name, addition)):
            folder_name = f"MIX{id_}_{mark or name}_{addition}"
        else:
            folder_name = folder_name or self.folder_name
        folder = self.root.joinpath(folder_name)
        folder.mkdir(exist_ok=True)
        return folder

    def generate_works_name(self, data: dict) -> str:
        return self.cleaner.filter_name(
            self.split.join(
                data[i] for i in self.name_format), inquire=False, default=str(
                time())[:10])

    def create_works_folder(
            self,
            root: Path,
            name: str,
            pass_=False) -> Path:
        if pass_:
            return root
        return root.joinpath(name) if self.folder_mode else root

    def stop_download_and_save(self):
        if self.__pool:
            self.__pool.shutdown(wait=True)  # 停止线程池中所有任务
            print("下载任务已停止")

            # 访问受保护成员 _unfinished_work
            for task in self.__pool._unfinished_work:
                url, temp, actual, show, id_ = task.args
                if self.is_exists(temp):
                    self.save_file(temp, actual)
                    print(f"{show} 已下载，并保存为 {actual.resolve()}")
                    self.blacklist.update_id(id_)
                else:
                    print(f"{show} 下载未完成，未保存文件")

    @staticmethod
    def save_file(temp: Path, actual: Path):
        move(temp.resolve(), actual.resolve())

    @staticmethod
    def delete_file(path: Path):
        path.unlink()
        print(f"文件 {path.name} 已删除")

    @staticmethod
    def statistics_count(count: SimpleNamespace):
        print(f"跳过视频作品 {len(count.skipped_video)} 个")
        print(f"跳过图集作品 {len(count.skipped_image)} 个")
        print(f"下载视频作品 {len(count.downloaded_video)} 个")
        print(f"下载图集作品 {len(count.downloaded_image)} 个")
